import pandas as pd
from scipy.stats import pearsonr, spearmanr

import itertools
import random

import hashlib

def generate_combinations(items):
    # 排序输入序列
    items.sort()

    # 设置随机种子
    random.seed(42)

    # 生成组合
    combs = list(itertools.combinations(items, 2))

    MOD = 10**9
    OFFSET = 2 * MOD  # 一个较大的偏移量，确保哈希值为正
    # 为每个组合分配唯一的编号
    secu_combs = {}
    for i, comb in enumerate(combs):
        comb_hash = hashlib.md5('_'.join(comb).encode()).hexdigest()
        secu_comb = (int(comb_hash, 16) + OFFSET) % MOD
        secu_combs[secu_comb] = {'code1': comb[0], 'code2': comb[1]}

    result = []
    for secu_comb, comb in secu_combs.items():
        result.append({f'{secu_comb}': comb})

    return result

def find_high_correlation_periods(data, date_col='trading_date', code_col='secu_code', value_col='change_pct', window_size=3):
    """
    找出数据中每个 code 与其他 code 最相关的时间段
    Args:
    data: 包含时序数据的DataFrame，包含日期列、代码列和值列
    date_col: 日期列的列名
    code_col: 代码列的列名
    value_col: 值列的列名
    window_size: 时间段窗口大小

    Returns:
    一个DataFrame，包含每个时间段的相关性信息
    """
    results = []

    # 获取数据中所有唯一的代码
    codes = data[code_col].unique()

    # 组合所有代码的两两组合
    code_combinations = generate_combinations(codes)

    # 遍历所有的代码组合
    for code_combs in code_combinations:
        for secu_comb, code_comb in code_combs.items():
            code_x, code_y = code_comb['code1'], code_comb['code2']
            # 获取第一组数据的值列
            values_x = data[data[code_col] == code_x][value_col].values

            # 获取第二组数据的值列
            values_y = data[data[code_col] == code_y][value_col].values

            # 计算每个窗口大小为5的时间段的相关系数
            for start_index in range(len(values_x) - window_size + 1):
                end_index = start_index + window_size

                # 使用pearsonr计算当前时间段的相关系数
                pearson_corr_value, _ = pearsonr(values_x[start_index:end_index], values_y[start_index:end_index])

                # 使用spearmanr计算当前时间段的相关系数
                spearman_corr_value, _ = spearmanr(values_x[start_index:end_index], values_y[start_index:end_index])

                corr_value = (pearson_corr_value + spearman_corr_value)/2

                # 保存相关性和对应的起始日期
                trading_date = data[data[code_col] == code_x][date_col].iloc[start_index]
                results.append({'secu_comb': secu_comb, 'trading_date': trading_date, 'code1': code_x, 'code2': code_y, 'corr_value': corr_value, 'pearson_corr_value': pearson_corr_value, 'spearman_corr_value': spearman_corr_value})

    # 将结果转换为DataFrame
    return pd.DataFrame(results)
#
# # 示例数据
# data = pd.DataFrame({
#     'date': pd.date_range(start='2022-01-01', periods=11).append(pd.date_range(start='2022-01-01', periods=11)).append(pd.date_range(start='2022-01-01', periods=11)),
#     'code': ['x']*11 + ['y']*11 + ['z']*11,
#     'value': [1, 90, 3, 4, 5, 6, 1245, 3, 55, 2, 4, 7, -210, 4, 5, 6, 7, 52, 3, 55, 2, 4, 1, 90, 3, 4, 5, 6, 1245, 3, 55, 2, 4]
# })
#
# # 设置时间段窗口大小
# window_size = 5
#
# # 找出相关性较高的时间段
# results_df = find_high_corr_valueelation_periods(data, window_size=window_size)
# print(results_df)



def process_combinations(df):
    # 过滤相关系数大于0.75的股票组合
    # df = df[df['corr_value'] >= 0.75].copy()

    # 将日期列转换为日期时间格式
    df['trading_date'] = pd.to_datetime(df['trading_date'])

    # 先按secu_comb和trading_date进行排序
    df = df.sort_values(by=['secu_comb', 'trading_date'])

    # 初始化 start_date 和 end_date 列
    df['start_date'] = df['trading_date']
    df['end_date'] = df['trading_date']

    # 遍历每个组合
    for secu_comb, comb_group in df.groupby('secu_comb'):
        # 获取该组合的第一个 trading_date
        first_trading_date = comb_group['trading_date'].iloc[0]
        prev_date = None  # 前一个日期
        prev_index = None  # 前一个索引

        update_flag = False
        # 遍历组合内的每一行
        for index, row in comb_group.iterrows():
            # 如果prev_date不为空，即不是第一行

            if prev_date is not None:
                # 当前日期和前一个日期的差值
                date_diff = (row['trading_date'] - prev_date).days

                # 如果差值大于等于10，更新start_date为当前日期
                if date_diff >= 10:
                    df.loc[index, 'start_date'] = row['trading_date']
                    update_flag = False
                    prev_date = None

            # 如果股票组合相关性第一次大于0.8，则开始更新start_date, end_date
            if not(update_flag) and row['corr_value'] >= 0.8:
                first_trading_date = row['trading_date']
                update_flag = True

            if update_flag:

                # 否则更新end_date为下一行的日期
                df.loc[index, 'start_date'] = first_trading_date
                df.loc[index, 'end_date'] = row['trading_date']

                # 更新prev_date和prev_index
                prev_date = row['trading_date']
                prev_index = index

        if prev_index:
            # 处理最后一行
            df.loc[prev_index, 'start_date'] = first_trading_date
            df.loc[prev_index, 'end_date'] = prev_date
            df['duration'] = (df['end_date'] - df['start_date']).dt.days

    return df

# 测试
# data = {
#     'secu_comb': [539261226]*7 + [66602460]*4,
#     'trading_date': ['2022-01-01', '2022-01-02', '2022-02-03', '2022-02-04', '2022-02-05', '2022-01-06', '2022-01-07', '2022-01-01', '2022-01-02', '2022-01-03', '2022-01-04'],
#     'code1': ['x']*11,
#     'code2': ['y']*7 + ['z']*4,
#     'corr_value': [-0.799687, -0.499583, 0.999332, 0.998822, 0.755565, 0.756078, 0.757136, 1.0, 1.0, 1.0, 1.0],
#     'pearson_corr_value': [-0.999374, -0.999166, 0.998664, 0.997645, 0.611129, 0.612156, 0.614271, 1.0, 1.0, 1.0, 1.0],
#     'spearman_corr_value': [-0.6, 0.0, 1.0, 1.0, 0.9, 0.9, 0.9, 1.0, 1.0, 1.0, 1.0]
# }
#
# df = pd.DataFrame(data)
# print(df)
# result_df = process_combinations(df)
# print(result_df)
#
